NBCE：使用朴素贝叶斯扩展LLM的Context处理长度

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

在 LLM 时代还玩朴素贝叶斯（Naive Bayes）？

这可能是许多读者在看到标题后的首个想法。确实如此，当古老的朴素贝叶斯与前沿的 LLM 相遇时，产生了令人惊讶的效果——我们可以直接扩展现有 LLM 模型的 Context 处理长度，无需对模型进行微调，也不依赖于模型架构，具有线性效率，而且效果看起来还不错——这就是本文所提出的 NBCE（Naive Bayes-based Context Extension）方法。

摸石过河

假设 为要生成的 token 序列， 是给定的若干个相对独立的 Context 集合（比如 个不同的段落，至少不是一个句子被分割为两个片段那种），假设它们的总长度已经超过了训练长度，而单个 加 还在训练长度内。我们需要根据 生成 ，即估计 。

简单来说，朴素贝叶斯就是“贝叶斯公式+独立假设”。根据贝叶斯公式：

这里的 ，是省去了与 无关的常数因子。根据（条件）独立假设：

所以有

再次根据贝叶斯公式 ，得到

或者

这里的 都可以直接用现有的 LLM 进行计算，而且只要是语言模型都行，跟架构无关，也不需要用长文本微调。其中， 是单个 Context 所预测的概率， 则无 Context（或者 Context 为空）的概率，并且多个 Context 可以放在同一个 batch 中并行计算，计算量随着 Context 数的增加是线性增长的。

抽丝剥茧

当然，朴素贝叶斯依赖于独立假设，这会限制它的实际效果。为了“青出于蓝而胜于蓝”，我们不妨将式（5）进一步“抽丝剥茧”、“去芜存菁”，以达到更好的效果。

首先我们记 ，以及

并设 ，那么式（5）可以重写为

重写为上述形式后，自然而言地引出了两个问题：

1. 如果将 作为超参数来调，是否可能取得更好的效果？

2.  就是 的 Average Pooling，那么换成其他 Pooling 方法（简记为 ）是否有更好的效果？即

于是笔者在 7B 模型上围绕这两个问题进行调试，得到的初步结论是：在阅读理解场景中 Max Pooling 配合 ，用 Greedy Search 总体表现比较好，然而 Random Sample 出来的结果基本不可读。

最终方案

为什么会出现 Greedy Search 好而 Random Sample 差的情况呢？我们知道，Random Sample 是“按照分布采样”，它的效果差说明 Max Pooling 的结果不是一个合理的分布；而 Greedy Search 只关心最大概率者，而不关心分布的合理性，它的效果好告诉我们概率最大的 token 正确性较高。

概率越大说明不确定性越低，所以为了改善 Random Sample 的效果，我们将 Pooling 方式改为直接输出不确定性最低的那个分布：

代入到式（8），就是最终的 NBCE（Naive Bayes-based Context Extension）。

值得指出的是，虽然我们的出发点是朴素贝叶斯，但一般化后的式（8）已经超出了常规的朴素贝叶斯的范畴，同时保留了朴素贝叶斯的可解释性。不难看出，式（8）的形式很是直观：

1、不同 Context 的预测结果通过方法 聚合（或者说投票）在一起（权重为 ），并减去无 Context 的预测结果（权重为 ）；

2、之所以要减去无 Context 预测结果，是为了让模型更加倾向于结合 Context 而不是纯粹根据自身知识储备来回答（注：3天后出现在 Arxiv 的论文《Trusting Your Evidence: Hallucinate Less with Context-aware Decoding》[1] 也提出了相同的技巧用来减少幻觉）；

3、不同场景可以选择不同的 ，比如需要结合 Context 做阅读理解的，可以考虑较大的 ，如果偏向于自由创作，则选择较小的 ，笔者认为 都是合理的。

参考实现

下面给出 NBCE 的参考实现：

Github: 

https://github.com/bojone/NBCE

从演示代码可以看出，NBCE 的实现很简单，只需要修改一下解码函数中的logits构建方式，跟解码算法的选择并不冲突。

▲ Naive Bayes-based Context Extension（NBCE）示意图

所给的 Demo 包含 12 段不同的 Context，总长度为 9000 多字，连同 8 个问题一次性输入到模型中（模型训练长度为 2048，参数量为 7B，可以在 OpenBuddy 下载），模型能够逐一根据所给 Context 正确回答这 8 个问题。值得指出的是，所有的 Context、问题和答案加起来，超过了 1 万字！另外，有朋友简单尝试了简历匹配和作文打分应用，效果也尚可，非常建议大家亲自调试一下。

相关工作

扩展 LLM 的 Context 长度其实已有不少，但多数是通过结合检索或者摘要的方式来缩短样本的长 Context，如 Unlimiformer [2]。由于不是直接处理长 Context，因此通常无法做精细的阅读理解，而且这些方案往往需要在训练阶段就考虑进去，而不是事后即插即用到已有的 LLM 模型中。

在 NBCE 之前，能够不微调地扩展 Context 长度的方案是 Parallel Context Window（下面简称 PCW），出自论文《Parallel Context Windows for Large Language Models》[3]和《Structured Prompting: Scaling In-Context Learning to 1,000 Examples》[4]，两篇论文是同一时期不同作者的工作，但所提的方法只有细微的差别，因此这里都将它们叫做 PCW。

PCW 适用于 Self Attention 模型，主要修改包括 Position Encoding 和 Attention Mask，如下图所示：

▲ Parallel Context Window

首先确定 Context 的最大长度 （图中为 6），然后每个 Context 的最后一个位置编码为 ，倒数第二个位置编码为 ，...，依此类推，这种编码方式我们称为“右对齐”（或者“左缩进”）；另一边，对于 Task Tokens 部分（Prompt+ 生成内容），我们的位置编码是 。

每个 Context 单独编码，所以对应的 Attention Mask 是分块对角矩阵，而因为是 LM，所以是分块对角下三角阵；至于 Task Tokens 部分需要结合所有的 Context，所以它需要 Attention 到所有 Context（以及它自身）。这样一来，如果将每个 Context 单独拿出来，和 Task Tokens 拼在一起，其 Attention 模式就跟原本的 LM 一致了。

或许有读者看出，其实 NBCE 跟 PCW 有着很相似的特性，比如对于 Context 都是无序的、平权的。事实上，如果将 NBCE 应用到单层单头注意力模型中，那么结果大致上就是 PCW。为了显示这一点，我们写出单层单头注意力的语言模型为

所以大致上有 ，接着代入到式（7）并取 ，得到

这里假设的是 是单个 token，但其实已经不失一般性了， 是拼接的意思。在上式中， 是作为一个连续片段来推理的（NBCE 的设定），所以它们的位置编码相邻，而 构成了 与所有 的一个整体 Attention（求和同样是 1），这些特性跟 PCW 其实是一致的，PCW 只不过是以 Attention Mask 的方式更优雅地整合到每一层中。

因此，PCW 大致上就是 Average Pooling 版的 NBCE，我们实测也发现它跟 Average Pooling 版的 NBCE 有着相似的缺点——当 Context 数据增加时，输出的结果开始不够准确，具体表现为主题相关，但是作为问题的答案来说是错误的。

延伸思考

NBCE 的一大缺点是无序性，即无法识别 Context 的输入顺序，这在续写故事等场景可能表现欠佳。为了缓解这一点，可以考虑在每一个 Context 前面加个能指示序信息的 prefix，就好比小说中的“第一章”、“第二章”那样。

总的来说，目前笔者关于 NBCE 的测试都限于“阅读理解”场景，即“理解”长文本，能否用此方法来“生成”长文本，还是个未知数，期待大家的测试结果。

此外，还有一个有意思的问题是：

既然朴素贝叶斯都能在 LLM 领域能派上用场，那么其他传统概率模型（比如HMM）是否也能在 LLM 领域有它们的一席之地呢？

补充说明和分析

上文中，我们介绍了一种基于朴素贝叶斯来扩展 LLM 的 Context 长度的方案 NBCE（Naive Bayes-based Context Extension）。由于它有着即插即用、模型无关、不用微调等优点，也获得了一些读者的认可，总的来说目前大家反馈的测试效果还算可以。

当然，部分读者在使用的时候也提出了一些问题。下文就结合读者的疑问和笔者的后续思考，对 NBCE 方法做一些补充说明和分析。

7.1 方法回顾

假设 T 为要生成的 token 序列， 是给定的若干个 Context，我们需要根据 生成 T，那么就需要估计 。根据朴素贝叶斯思想，我们得到

其中 ，细节请参考上文。NBCE 做了两点改动：1）将 作为超参数来调；2）将 换成一般的 Pooling 方法 。结果变为

最后，NBCE 选取的 Pooling 方案为“取最小熵者”：

7.2 截断预测

式 (12) 是标准的朴素贝叶斯结果，但笔者按照它实现时，发现随着 n 的增大，式 (12) 的效果慢慢变差直至完全乱码，所以反复调整后，最终选取“取最小熵者”作为 NBCE 的 Pooling 方案。但后来仔细一想，式 (12) 的这个表现是不正常的，因为朴素贝叶斯的唯一假设是 Context 之间相互独立，而我所测试的 Context 是随机找的若干篇新闻，一定程度上满足这个假设的，所以式 (12) 再怎么差也不至于完全乱码才对。

百思不得其解之际，微信上@孔某人提醒：语言模型的训练标签都是 One-Hot，所以预测结果除了头部（概率最高的部分）外其实都是不可信的。这个提示可谓是一针见血，让答案瞬间水落石出了：由于式 (12) 包含了 这一项，它会放大尾部的预测结果，如果尾部的预测结果不可靠，那么这个放大作用甚至会完全打乱头部的准确结果。

为什么它不会影响“取最小熵者”呢？因为最小熵的结果头部概率往往更大，尾部概率更小，所以即便 这一项放大了尾部，也依然无法胜过头部；而对于式 (12) 来说，它是所有预测结果的平均，则会弱化头部，以至于加上 后尾部胜过了头部。

有了这个线索，解决方法也就很明显了，给每个预测结果都加上 Top-P 或者 Top-K 截断就行，在 Github 的代码中，笔者选的是 Top-P 截断。

7.3 无穷处理

然而，事情还没完，经过截断后， 的尾部都变成了 ，这时候式 (12) 或式 (13) 都有可能出现 的无意义运算。总的来说，有下面几种情况：

其中，“情况一”和“情况三”能正常运算，“情况二”虽然也能正常运算，但其结果正无穷是不合理的，“情况四”则是非良定的无意义运算，也就是说，我们需要想办法修正“情况二”和“情况四”，这两种情况正好对应 ，所以我们将式 (13) 改为：

经过上述处理后，标准的朴素贝叶斯式 (12) 也能输出正常结果了（虽然最终效果还是不如取最小熵，但至少不会乱码了），并且修改后的代码对 Pooling 方式和 都更加鲁棒。

7.4 转移概率

当用于回答一些观点型问题或者偏向自由创作的问题时，NBCE 可能会出现在 Context 中反复跳转的问题。具体来说，就是模型并没有自信地注意到某个 Context 中，所以 之间相差不大，于是式 (14) 中的 结果不稳定，每步生成都选取到了不同的 Context 上，导致生成结果语义上的不连续甚至完全和 Context 不相关，加剧 LLM 的“幻觉”现象。

为了缓解这个问题，我们可以模仿转移概率的概念，适当给上一步所选的 Context 加权，让模型“非必要不跳转”，具体做法是引入参数 ，将式 (14) 改为

其中 k' 是上一步生成时选取的 Context 编号。这样一来，只有 时才会发生 Context 跳转，降低了跳转概率。

以上提到的修改都已经同步到 Github 中：

https://github.com/bojone/NBCE

7.5 适用场景

由于朴素贝叶斯所做的独立假设，不少读者可能会怀疑：当 Context 之间有明显的语义重叠时，NBCE 的效果会不会明显下降？或者说，NBCE 的适用场景是什么？

事实上，受限于独立假设的是标准的朴素贝叶斯，也就是式 (12)，一般化后的式 (13) 和式 (14) 已经基本上不受限于独立假设了。事实上，“取最小熵者”版本的 NBCE 实质就是用 LLM 的熵作为相似度对 Context 进行检索，每步生成都会更新检索结果。所以，NBCE 的适用场景是：假设要预测的答案可以分为若干个片段，每个片段只依赖于一个 Context。

基于这个结论，当我们只有一段长文本作为 Context（比如一篇小说）时，我们可以通过有重叠的滑窗自动将长 Context 划分为多个短 Context，而不一定需要人工分割为相对独立的片段，因为刚才的结论告诉我们，NBCE 的适用场景跟 Context 的重叠性无关。至于为什么重叠地滑窗，则只是为了能尽可能依赖单个 Context 就能输出完整的结果。

以下两种场景用 NBCE 效果大概率是不会好的：

1. 有序的 Context：这是指当生成结果强烈依赖于 Context 的输入顺序（或者更复杂的嵌套结构）时，NBCE 通常效果不好，因为 NBCE 保留了朴素贝叶斯的无序性，这种场景的典型例子是给小说写摘要（小说被切割为多个 Context），一个治标不治本的方案是人工给每段 Context 加上可以识别顺序的标记，如“第 xx 章”；

2. 耦合的 Context：这是指必须要结合两个或更多 Context 构建输出时，NBCE 效果不好，因为 NBCE 每次只选了一个 Context，@孔某人给出了一个典型例子“已知 x > 1，且 x < 0，求 x 解集”，假设两个条件被分为两个 Context，那么必须结合两个 Context 才能输出正确答案“空集”，只看单个 Context 无法确定是空集。

如果 NBCE 还要继续做下去的话，大体上也是围绕以上两个场景进行改进。

文章小结

本文提出了 NBCE（Naive Bayes-based Context Extension），它基于朴素贝叶斯思想来扩展 LLM 的 Context 处理长度，有着即插即用、模型无关、无须微调、线性效率、实现简单等优点，并且看上去效果还不错，欢迎大家测试。

此外，本文还补充介绍了 Context 长度扩展方案 NBCE 的一些后续更新及分析，并进一步讨论了 NBCE 的适用场景。

参考文献

[1] https://arxiv.org/abs/2305.14739

[2] https://arxiv.org/abs/2305.01625

[3] https://arxiv.org/abs/2212.10947

[4] https://arxiv.org/abs/2212.06713

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·